package jast;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Stack;

import org.antlr.v4.runtime.ParserRuleContext;
import org.antlr.v4.runtime.tree.ParseTreeProperty;
import org.antlr.v4.runtime.tree.TerminalNodeImpl;

import parser.JavaParser;
import path.BranchNode;
import path.SwitchBranchNode;

public class CfgListener extends JavaBaseListener {
	JavaParser parser;
	//Node Index
	private int index;
	//Exit Node
	BranchNode methodExit;
	BranchNode currentNode;
	ArrayList<BranchNode> methodStart = new ArrayList<BranchNode>();
	ParseTreeProperty<BranchNode> values = new ParseTreeProperty<BranchNode>();
	HashMap<BranchNode,Integer> position = new HashMap<BranchNode,Integer>();
	HashMap<String,ParserRuleContext> labelMap = new HashMap<String,ParserRuleContext>();
	Stack<BranchNode> stack = new Stack<BranchNode>();
	public CfgListener(JavaParser parser){
		this.parser = parser;
	}
	/**
	 * PackageDeclaration
	 */
	public void enterPackageDeclaration(JavaParser.PackageDeclarationContext ctx) { 
//		System.out.println("Package:"+ctx.getText());
	}
	
	public void enterClassDeclaration(JavaParser.ClassDeclarationContext ctx) {
//		System.out.println(ctx.Identifier()+" {");
	}
	
	public void exitClassDeclaration(JavaParser.ClassDeclarationContext ctx) {
//		System.out.println("}");
	}
	
	public void enterMethodDeclaration(JavaParser.MethodDeclarationContext ctx) { 
		BranchNode start = new BranchNode("","",-1,0);
		BranchNode exit = new BranchNode("","",-2,0);
		methodStart.add(start);
		currentNode = start;
		methodExit = exit;
		index = 0;
	}
	
	public void exitMethodDeclaration(JavaParser.MethodDeclarationContext ctx) { 
		currentNode.addLeft(methodExit);
	}
	
	public void enterParExpression(JavaParser.ParExpressionContext ctx) {
//		System.out.println(ctx.getText());
//		System.out.println(ctx.getChildCount());
		for(int i=0;i<ctx.getChildCount();i++){
			//ParseTree
//			System.out.println(ctx.children.get(i).toStringTree(parser));
//			System.out.println(ctx.children.get(i) instanceof JavaParser.ExpressionContext);
		}
//		System.out.println(ctx.toStringTree(parser));
		
//		System.out.println("Expressions:");
//		for(int i=0;i<ctx.getChildCount();i++){
//			if(ctx.children.get(i) instanceof JavaParser.ExpressionContext)
//				System.out.println(ctx.children.get(i).getText());
//		}
	}
	
	public void enterSwitchBlockStatementGroup(JavaParser.SwitchBlockStatementGroupContext ctx) {
		BranchNode parentNode = stack.peek();//当前所处位置
		if(parentNode.type == 3){
			if(position.containsKey(parentNode))
				position.put(parentNode, position.get(parentNode)+1);//跳转到下一个case
			else
				position.put(parentNode, 0);
			currentNode = ((SwitchBranchNode)parentNode).caseList.get(position.get(parentNode));
		}
	}
	
	public void exitSwitchBlockStatementGroup(JavaParser.SwitchBlockStatementGroupContext ctx) {
		if(currentNode.getLeft() == null){
			BranchNode parentNode = stack.peek();
			if(position.get(parentNode)+1 == ((SwitchBranchNode)parentNode).caseList.size())
				currentNode.addLeft(((SwitchBranchNode)parentNode).exitNode);
			else{
				BranchNode nextCaseNode = ((SwitchBranchNode)parentNode).caseList.get(position.get(parentNode)+1);
				currentNode.addLeft(nextCaseNode);
			}
		}
	}
	
	public void enterStatement(JavaParser.StatementContext ctx) {
		String firstLabel = ctx.children.get(0).getText();
		if(firstLabel.equals("if")){
			BranchNode parentNode = new BranchNode("","",-1,0);
			BranchNode node = new BranchNode("","",++index,0);
//			System.out.println(currentNode.elseIndex+" "+ctx.start.getTokenIndex());
			updateCurrentNode(ctx);
			//创建父节点
			currentNode.addLeft(parentNode);
			currentNode = parentNode;
			currentNode.addLeft(node);
			//if-else结束节点
			BranchNode exitNode = new BranchNode("","",-1,0);
			currentNode.exitNode = exitNode;
			stack.push(currentNode);//进入if-else内部
			values.put(ctx,currentNode);
//			System.out.println(firstLabel);
//			System.out.println("Enter"+index);
			if(ctx.children.size()>3&&ctx.children.get(3).getText().equals("else")){
				BranchNode elseNode = new BranchNode("","",index,1);
				currentNode.addRight(elseNode);
				//if块最后一个Token的位置
				currentNode.elseIndex = ((JavaParser.StatementContext)ctx.children.get(2)).stop.getTokenIndex();
			}
			else
				//无else，直接以if块作为当前节点
				currentNode = node;
		}
		else if(firstLabel.equals("for")){
			updateCurrentNode(ctx);
			//不在内层嵌套的情况，直接以当前节点做父节点
			//判断
			BranchNode forNode = new BranchNode("","",-1,0);
			//与while类似
			forNode.type = 2;
			//循环体节点
			BranchNode bodyNode = new BranchNode("","",++index,0);
			currentNode.addLeft(forNode);
			forNode.addLeft(bodyNode);
			//出节点
			BranchNode exitNode = new BranchNode("","",-1,0);
			forNode.addRight(exitNode);
			forNode.exitNode = exitNode;
			stack.push(forNode);//进入for内部
			if(labelMap.containsValue(ctx))
				values.put(ctx,forNode);
			currentNode = bodyNode;
//			System.err.println(firstLabel);
		}
		else if(firstLabel.equals("while")){
			updateCurrentNode(ctx);
			//不在内层嵌套的情况，直接以当前节点做父节点
			//判断节点
			BranchNode whileNode = new BranchNode("","",-1,0);
			whileNode.type = 2;
			//循环体节点
			BranchNode bodyNode = new BranchNode("","",++index,0);
			//出节点
			BranchNode exitNode = new BranchNode("","",-1,0);
			currentNode.addLeft(whileNode);
			whileNode.addLeft(bodyNode);
			whileNode.addRight(exitNode);
			whileNode.exitNode = exitNode;
			stack.push(whileNode);//进入while内部
			if(labelMap.containsValue(ctx))
				values.put(ctx,whileNode);
			currentNode = bodyNode;
//			System.out.println(ctx.children.get(1).getText()+" "+index);
		}
		else if(firstLabel.equals("do")){
			updateCurrentNode(ctx);
			//创建循环体节点
			BranchNode bodyNode = new BranchNode("","",++index,0);
			bodyNode.type = 1;
			currentNode.addLeft(bodyNode);
			//创建退出虚节点，避免以while节点作退出节点会发生冲突（右节点已使用，且不能改变）
			BranchNode exitNode = new BranchNode("","",-1,0);
			bodyNode.exitNode = exitNode;
			stack.push(bodyNode);//进入do-while内部
			if(labelMap.containsValue(ctx))
				values.put(ctx,bodyNode);
			currentNode = bodyNode;
//			System.err.println("Enter do-while:"+ctx.children.get(3).getText());
		}
		else if(firstLabel.equals("switch")){
//			System.out.println("ENTER SWITCH:"+ctx.children.get(1).getText());
			updateCurrentNode(ctx);
			SwitchBranchNode switchNode = new SwitchBranchNode("","",-1,0);
			currentNode.addLeft(switchNode);
//			for(int i=2;i<ctx.children.size();i++)
//				System.out.println(i+":"+ctx.children.get(i).getText());
			index++;
			//case
			for(int i=3;i<ctx.children.size()-1;i++){
				BranchNode bodyNode = new BranchNode("","",index,i-3);
				switchNode.caseList.add(bodyNode);
			}
			BranchNode exitNode = new BranchNode("","",-1,0);
			switchNode.exitNode = exitNode;
			stack.push(switchNode);//进入switch-case内部
			if(labelMap.containsValue(ctx))
				values.put(ctx,switchNode);
		}
		else if(firstLabel.equals("break")){
			updateCurrentNode(ctx);
			//是否有Identifier
			if(ctx.children.size() == 2){
				BranchNode parentNode = stack.peek();
				//寻找第一个switch-case,while,do-while,for节点
				for(int i=0;i<stack.size();i++){
					parentNode = stack.get(i);
					if(parentNode.type > 0)
						break;
				}
				//将currentNode连接结束节点
				currentNode.addLeft(parentNode.exitNode);
				//如果该节点在if-else中，则在退出if-else时不再对该节点进行处理
				currentNode.type = -1;
			}
			else{
				String label = ctx.children.get(1).getText();
				ParserRuleContext prCtx = labelMap.get(label);
				BranchNode bNode = values.get(prCtx);
				//连接循环结束点
				currentNode.addLeft(bNode.exitNode);
				//如果该节点在if-else中，则在退出if-else时不再对该节点进行处理
				currentNode.type = -1;
			}
		}
		else if(firstLabel.equals("continue")){
			updateCurrentNode(ctx);
			//是否有Identifier
			if(ctx.children.size() == 2){
				BranchNode parentNode = stack.peek();
				//寻找第一个while,do-while,for节点
				for(int i=0;i<stack.size();i++){
					parentNode = stack.get(i);
					if(parentNode.type > 0)
						break;
				}
				//将currentNode连接开始节点
				currentNode.addLeft(parentNode);
				//如果该节点在if-else中，则在退出if-else时不再对该节点进行处理
				currentNode.type = -1;
			}
			else{
				String label = ctx.children.get(1).getText();
				ParserRuleContext prCtx = labelMap.get(label);
				BranchNode bNode = values.get(prCtx);
				//连接循环开始点
				currentNode.addLeft(bNode);
				//如果该节点在if-else中，则在退出if-else时不再对该节点进行处理
				currentNode.type = -1;
			}
		}
		else if(firstLabel.equals("return")){
			updateCurrentNode(ctx);
			//连接函数结束节点
			currentNode.addLeft(methodExit);
			currentNode.type = -1;
		}
		else if(ctx.children.get(0) instanceof TerminalNodeImpl){
			if(firstLabel.equals("try")||firstLabel.equals("synchronized")||firstLabel.equals("assert")||firstLabel.equals("throw")||firstLabel.equals(";"))
				return;
			//Identifier ':' statement
			//判断后续是否是do，while，for
			JavaParser.StatementContext statementContext = (JavaParser.StatementContext)ctx.children.get(2);
			String nextLabel = statementContext.children.get(0).getText();
			if(nextLabel.equals("for")||nextLabel.equals("while")||nextLabel.equals("do")){
				//记录Identifier对应的循环
				labelMap.put(firstLabel, statementContext);
			}
			else{
//				System.err.println("Label（"+firstLabel+"）应在循环前！");
			}
		}
	}
	/**
	 * 更新当前节点
	 * @param ctx
	 */
	private void updateCurrentNode(JavaParser.StatementContext ctx) {
		if(!stack.isEmpty()){
			BranchNode parentNode = stack.peek();//外层嵌套
			//判断节点在if块内还是else块内
			if(parentNode.elseIndex>0&&ctx.start.getTokenIndex()>parentNode.elseIndex){
				currentNode = parentNode.getRight();
				while(currentNode.getLeft()!=null)
					currentNode = currentNode.getLeft();
			}
			else if(parentNode.elseIndex != -1){
				currentNode = parentNode.getLeft();
				while(currentNode.getLeft()!=null)
					currentNode = currentNode.getLeft();
			}
		}
	}
	
	public void exitStatement(JavaParser.StatementContext ctx) {
		if(ctx.children.get(0).getText().equals("if")){
			BranchNode parentNode = values.get(ctx);
			
			BranchNode lastLeft = parentNode.getLeft();
			while(lastLeft.getLeft()!=null&&lastLeft.type != -1){
				if(lastLeft.type == 2)
					lastLeft = lastLeft.getRight();
				else
					lastLeft = lastLeft.getLeft();
			}
			if(lastLeft.type != -1)
				lastLeft.addLeft(parentNode.exitNode);
			if(parentNode.getRight()!=null){
				lastLeft = parentNode.getRight();
				while(lastLeft.getLeft()!=null&&lastLeft.type != -1){
					if(lastLeft.type == 2)
						lastLeft = lastLeft.getRight();
					else
						lastLeft = lastLeft.getLeft();
				}
				if(lastLeft.type != -1)
					lastLeft.addLeft(parentNode.exitNode);//else
			}
			else{
				parentNode.addRight(parentNode.exitNode);
			}
			stack.pop();//退出if-else内部
			//不会成为孤立节点，因为若if和else均退出，下面的语句也就不会执行，成为孤立代码
			currentNode = parentNode.exitNode;
//			System.out.println("EXIT(ELSE):"+parentNode.getLeft().branchNum);
		}
		else if(ctx.children.get(0).getText().equals("while")){
			BranchNode whileNode = stack.pop();	
			//寻找循环体
			BranchNode bodyNode = whileNode.getLeft();
			bodyNode = getLastNode(bodyNode);
			bodyNode.addRight(whileNode);

			currentNode = whileNode.getRight();//退出while内部
		}
		else if(ctx.children.get(0).getText().equals("do")){
			BranchNode bodyNode = stack.pop();
			//寻找循环体结束节点
			BranchNode bodyEndNode = bodyNode;
			//避免该节点是最后一个节点，进入下面的方法会得到bodyNode.exitNode，使得部分节点悬空
			if(bodyEndNode.getLeft() != null){
				bodyEndNode = bodyEndNode.getLeft();
				bodyEndNode = getLastNode(bodyEndNode);
			}
			BranchNode whileNode = new BranchNode("","",-1,0);
			bodyEndNode.addLeft(whileNode);
			//while节点右节点连接循环体，在其他位置寻找结束节点时就不用单独考虑，不会进入循环
			whileNode.addRight(bodyNode);
			whileNode.addLeft(bodyNode.exitNode);
			//设置当前节点为退出节点
			currentNode = bodyNode.exitNode;
		}
		else if(ctx.children.get(0).getText().equals("for")){
			BranchNode forNode = stack.pop();	
			//寻找循环体
			BranchNode bodyNode = forNode.getLeft();
			bodyNode = getLastNode(bodyNode);
			bodyNode.addRight(forNode);
			//退出while内部
			currentNode = forNode.getRight();			
		}
		else if(ctx.children.get(0).getText().equals("switch")){
			SwitchBranchNode switchNode = (SwitchBranchNode) stack.pop();
			//退出switch-case内部
			currentNode = switchNode.exitNode;
		}
	}
	/**
	 * 得到后续的最后一个节点
	 * @param bodyNode
	 * @return
	 */
	private BranchNode getLastNode(BranchNode bodyNode) {
		while(bodyNode.getLeft()!=null){
			if(bodyNode.exitNode != null)
				bodyNode = bodyNode.exitNode;
			else
				bodyNode = bodyNode.getLeft();
		}
		return bodyNode;
	}
}